# 倒立摆系统的LQR控制,例子:xd = [0,0]
#  1.两种方式计算反馈增益F，并且对两种方式计算的F值进行使用
# 方法1是在程序中硬算，方法2是调用<F1_LQR_Gain>函数
# 2.调整输入的阶数，1阶和2阶时候的使用，调整B，u0，R

import F_LQR_Gain 
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import time
import scipy.linalg as la
from scipy.signal import StateSpace

# 1.定义系统
g = 10
d = 1

# x[k+1] = Ax[k] + Bu[k]
A = np.array([[0,1],[g/d,0]])
B = np.array([[0],[1]])
C = np.array([[1,0]])
D = 0
# 创建状态空间模型
# sys = StateSpace(A, B, C, D)
Ts = 0.1
sys_continuous = ct.StateSpace(A, B, C, D)
sys_discrete = ct.sample_system(sys_continuous, Ts, method='zoh')
# 获取行数
n = len(A)
# 获取列数
p =  B.shape[1]
# 离散化后的状态矩阵和输入矩阵
A = sys_discrete.A
B = sys_discrete.B
# 2.初始化系统
x0 = np.array([[1],[0]]).reshape(1,-1)
x = x0
# 输入
u0 = np.array([[0]]).reshape(1,-1)
u = u0
# 定义系统运行步数
k_steps = 100
x_history = np.zeros((n,k_steps+1))
x_history[:,0] = x
u_history = np.zeros((p,k_steps))
u_history[:,0] = u
# print(x,"\n",x.shape)
# print(x_history,"\n",x_history.shape)
# 设置权重
Q = np.array([[1,0],[0,1]])
S = np.array([[1,0],[0,1]])
R=0.1
N = k_steps
P_k = S
# 反馈gain(F(N-k))
# 3.计算反馈增益
# 3.1.1方法1：在程序中硬算
# for k in range(0,N):
#     F = np.linalg.inv(R +B.T@P_k@B)@B.T@P_k@A
#     P_k = (A-B@F).T@P_k@(A-B@F) + F.T*R*F + Q
#     if k == 0 :
#         F_N = F
#     else :
#         F_N = np.vstack((F,F_N))


# 3.2.1方法2：调用<F1_LQR_Gain>函数
F = F_LQR_Gain.LQR_Gain(A,B,Q,R,S)

for k in range(1,k_steps+1):
    # # 方法3.1.2
    # u = -F_N[(k-1)*p:k*p,:]@x_history[:,k-1].reshape(-1,1)
    # 方法3.2.2
    u = -F@x_history[:,k-1].reshape(-1,1)
    # 3.3 循环更新x，记录x_history、u_history，方法1/2此部分程序相同
    x=A@x_history[:,k-1].reshape(-1,1)+B@u
    x_history[:,k]=x.T
    u_history[:,k-1]=u
    

"""
# # 写入数据到表格中
# from openpyxl import Workbook
# # 创建 Excel 工作簿
# wb = Workbook()
# ws = wb.active
# # 将矩阵数据写入工作表
# for row in x_history:
#     ws.append(row.tolist())
# # 保存 Excel 文件
# wb.save('x_history.xlsx')
# wb = Workbook()
# ws = wb.active
# # 将矩阵数据写入工作表
# for row in u_history:
#     ws.append(row.tolist())
# # 保存 Excel 文件
# wb.save('u_history.xlsx')

# # 创建 Excel 工作簿
# wb = Workbook()
# ws = wb.active
# # 将矩阵数据写入工作表
# for row in F_N:
#     ws.append(row.tolist())
# # 保存 Excel 文件
# wb.save('F_N.xlsx')
"""
# 4.制图
plt.figure()
plt.subplot(2,1,1)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
for i in range(0,n):
    plt.plot(x_history[i,:])
plt.subplot(2,1,2)

for i in range(0,p):
    plt.plot(u_history[i,:])

plt.ylabel('Input')
plt.grid()
plt.show()